package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
)

func handleClientConnection(clientConn net.Conn) {
	defer clientConn.Close()

	clientReader := bufio.NewReader(clientConn)

	// 读取客户端请求的第一行
	requestLine, err := clientReader.ReadString('\n')
	if err != nil {
		log.Printf("Error reading request line: %v", err)
		return
	}

	// 解析请求行
	parts := strings.Fields(requestLine)
	if len(parts) < 3 {
		log.Printf("Invalid request line: %s", requestLine)
		return
	}

	method, hostPort, protocol := parts[0], parts[1], parts[2]
	log.Printf("Received request: %s %s %s", method, hostPort, protocol)

	if method == http.MethodConnect {
		// 处理 HTTPS 请求
		handleHTTPS(clientConn, hostPort, clientReader)
	} else {
		// 处理 HTTP 请求
		handleHTTP(clientConn, requestLine, clientReader)
	}
}

func handleHTTPS(clientConn net.Conn, hostPort string, clientReader *bufio.Reader) {
	// 连接目标服务器
	targetConn, err := net.Dial("tcp", hostPort)
	if err != nil {
		log.Printf("Error connecting to target: %v", err)
		return
	}
	defer targetConn.Close()

	// 响应客户端连接已建立
	_, err = fmt.Fprintf(clientConn, "HTTP/1.1 200 Connection Established\r\n\r\n")
	if err != nil {
		log.Printf("Error writing response to client: %v", err)
		return
	}

	// 开始转发数据
	go transferData(targetConn, clientConn)
	transferData(clientConn, targetConn)
}

func handleHTTP(clientConn net.Conn, requestLine string, clientReader *bufio.Reader) {
	// 读取完整的HTTP请求
	var headers []string
	headers = append(headers, requestLine)
	for {
		line, err := clientReader.ReadString('\n')
		if err != nil {
			log.Printf("Error reading HTTP headers: %v", err)
			return
		}
		headers = append(headers, line)
		if line == "\r\n" {
			break
		}
	}

	// 解析目标地址
	request, err := http.ReadRequest(bufio.NewReader(strings.NewReader(strings.Join(headers, ""))))
	if err != nil {
		log.Printf("Error parsing HTTP request: %v", err)
		return
	}
	targetURL := request.URL
	if !strings.Contains(targetURL.Host, ":") {
		targetURL.Host += ":80"
	}

	// 连接目标服务器
	targetConn, err := net.Dial("tcp", targetURL.Host)
	if err != nil {
		log.Printf("Error connecting to target: %v", err)
		return
	}
	defer targetConn.Close()

	// 将HTTP请求写入目标服务器
	_, err = targetConn.Write([]byte(strings.Join(headers, "")))
	if err != nil {
		log.Printf("Error writing HTTP request to target: %v", err)
		return
	}

	// 开始转发数据
	go transferData(targetConn, clientConn)
	transferData(clientConn, targetConn)
}

// transferData 传递数据
func transferData(destination io.WriteCloser, source io.ReadCloser) {
	defer destination.Close()
	defer source.Close()
	io.Copy(destination, source)
}

func main() {
	// 定义命令行参数
	port := flag.String("port", "8000", "Port to listen on")
	flag.Parse()

	// 启动TCP服务器监听指定端口
	listener, err := net.Listen("tcp", ":"+*port)
	if err != nil {
		log.Fatalf("Failed to start TCP server: %v", err)
	}
	defer listener.Close()

	log.Printf("TCP proxy server is running on port %s", *port)

	for {
		clientConn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept connection: %v", err)
			continue
		}

		go handleClientConnection(clientConn)
	}
}
